import os
import subprocess
import tempfile
from typing import Optional, Tuple, Union
from fastapi import FastAPI, HTTPException, File, UploadFile, Response
from pydantic import BaseModel
import cv2
import numpy as np
import io
from starlette.responses import StreamingResponse

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# 设置ADB路径
ADB_PATH = os.path.join(os.getcwd(), 'platform-tools', 'adb')

def start_adb_server():
    """启动ADB服务"""
    try:
        result = subprocess.run(
            [ADB_PATH, 'start-server'],
            capture_output=True,
            text=True,
            check=True,
            timeout=30
        )
        print("ADB服务已启动")
        return result.stdout, result.stderr
    except subprocess.CalledProcessError as e:
        print(f"启动ADB服务失败: {e.stderr}")
    except subprocess.TimeoutExpired:
        print("启动ADB服务超时")

class DeviceConnectionError(Exception):
    """设备连接异常"""
    pass

class CommandExecutionError(Exception):
    """命令执行异常"""
    pass

class ADBController:
    """ADB控制器，封装ADB命令"""
    
    def __init__(self, device_id: Optional[str] = None):
        self.device_id = device_id
        self._connected = False
        
    def _run_adb_cmd(self, command: list) -> Tuple[str, str]:
        """执行ADB命令并返回标准输出和错误输出"""
        full_cmd = [ADB_PATH]
        if self.device_id:
            full_cmd.extend(['-s', self.device_id])
        full_cmd.extend(command)
        
        try:
            result = subprocess.run(
                full_cmd,
                capture_output=True,
                text=True,
                check=True,
                timeout=30
            )
            return result.stdout, result.stderr
        except subprocess.CalledProcessError as e:
            raise CommandExecutionError(f"ADB命令执行失败: {e.stderr}") from e
        except subprocess.TimeoutExpired:
            raise CommandExecutionError("ADB命令执行超时")
    
    def connect(self, device_id: str) -> str:
        """连接到设备"""
        self.device_id = device_id
        stdout, stderr = self._run_adb_cmd(['connect', device_id])
        
        if "unable to connect" in stderr.lower() or "failed" in stderr.lower():
            raise DeviceConnectionError(f"无法连接到设备 {device_id}: {stderr}")
            
        self._connected = True
        return stdout
    
    def disconnect(self) -> str:
        """断开设备连接"""
        if not self._connected or not self.device_id:
            return "设备未连接"
            
        stdout, stderr = self._run_adb_cmd(['disconnect', self.device_id])
        
        if stderr and "not connected" not in stderr.lower():
            raise DeviceConnectionError(f"断开连接失败: {stderr}")
            
        self._connected = False
        self.device_id = None
        return stdout
    
    def is_connected(self) -> bool:
        """检查设备是否已连接"""
        if not self.device_id:
            return False
            
        try:
            stdout, _ = self._run_adb_cmd(['get-state'])
            return 'device' in stdout.strip().lower()
        except CommandExecutionError:
            return False
    
    def take_screenshot(self) -> np.ndarray:
        """获取设备截图"""
        if not self._connected:
            raise DeviceConnectionError("设备未连接")
            
        try:
            result = subprocess.run(
                [ADB_PATH, '-s', self.device_id, 'exec-out', 'screencap -p'],
                capture_output=True,
                check=True
            )
            img_array = np.frombuffer(result.stdout, np.uint8)
            img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            
            if img is None:
                raise CommandExecutionError("无法解析截图数据")
                
            return img
        except Exception as e:
            raise CommandExecutionError(f"获取截图失败: {str(e)}") from e
    
    def tap(self, x: Union[int, float], y: Union[int, float]) -> str:
        """点击屏幕指定位置"""
        if not self._connected:
            raise DeviceConnectionError("设备未连接")
            
        stdout, stderr = self._run_adb_cmd(
            ['shell', 'input', 'tap', str(int(x)), str(int(y))]
        )
        
        if stderr:
            raise CommandExecutionError(f"点击失败: {stderr}")
            
        return stdout
    
    def swipe(self, x1: Union[int, float], y1: Union[int, float], 
              x2: Union[int, float], y2: Union[int, float], 
              duration: int = 500) -> str:
        """在屏幕上滑动"""
        if not self._connected:
            raise DeviceConnectionError("设备未连接")
            
        duration = max(0, int(duration))
        
        stdout, stderr = self._run_adb_cmd([
            'shell', 'input', 'swipe',
            str(int(x1)), str(int(y1)),
            str(int(x2)), str(int(y2)),
            str(duration)
        ])
        
        if stderr:
            raise CommandExecutionError(f"滑动失败: {stderr}")
            
        return stdout

# 创建FastAPI应用
app = FastAPI(title="ADB HTTP服务", description="通过HTTP接口控制Android设备")

# 设备连接状态管理
devices = {}

@app.get("/")
def read_root():
    """API根路径，返回欢迎信息"""
    return {"message": "欢迎使用ADB HTTP服务"}

@app.post("/devices/{device_id}/connect")
def connect_device(device_id: str):
    """连接到指定设备"""
    try:
        if device_id in devices:
            controller = devices[device_id]
            if controller.is_connected():
                return {"status": "success", "message": f"设备 {device_id} 已连接"}
        else:
            controller = ADBController()
            devices[device_id] = controller
            
        result = controller.connect(device_id)
        return {"status": "success", "message": result}
    except (DeviceConnectionError, CommandExecutionError) as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/devices/{device_id}/disconnect")
def disconnect_device(device_id: str):
    """断开指定设备的连接"""
    try:
        if device_id not in devices:
            return {"status": "success", "message": f"设备 {device_id} 未连接"}
            
        controller = devices[device_id]
        result = controller.disconnect()
        del devices[device_id]
        return {"status": "success", "message": result}
    except (DeviceConnectionError, CommandExecutionError) as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/devices/{device_id}/screenshot")
def get_screenshot(device_id: str):
    """获取设备截图"""
    try:
        if device_id not in devices:
            raise HTTPException(status_code=404, detail=f"设备 {device_id} 未连接")
            
        controller = devices[device_id]
        if not controller.is_connected():
            raise HTTPException(status_code=400, detail=f"设备 {device_id} 未连接或已断开")
            
        img = controller.take_screenshot()
        _, img_encoded = cv2.imencode('.png', img)
        img_bytes = img_encoded.tobytes()
        
        return Response(content=img_bytes, media_type="image/png")
    except CommandExecutionError as e:
        raise HTTPException(status_code=400, detail=str(e))




class TapData(BaseModel):
    x: float
    y: float

@app.post("/devices/{device_id}/tap")
def tap_screen(device_id: str, data: TapData):
    """在设备屏幕上点击指定位置"""
    try:
        if device_id not in devices:
            raise HTTPException(status_code=404, detail=f"设备 {device_id} 未连接")
            
        controller = devices[device_id]
        if not controller.is_connected():
            raise HTTPException(status_code=400, detail=f"设备 {device_id} 未连接或已断开")
            
        result = controller.tap(data.x, data.y)
        return {"status": "success", "message": result}
    except (DeviceConnectionError, CommandExecutionError) as e:
        raise HTTPException(status_code=400, detail=str(e))




class SwipeData(BaseModel):
    x1: float
    y1: float
    x2: float
    y2: float
    duration: int = 500

@app.post("/devices/{device_id}/swipe")
def swipe_screen(device_id: str, data: SwipeData):
    """在设备屏幕上滑动"""
    try:
        if device_id not in devices:
            raise HTTPException(status_code=404, detail=f"设备 {device_id} 未连接")
            
        controller = devices[device_id]
        if not controller.is_connected():
            raise HTTPException(status_code=400, detail=f"设备 {device_id} 未连接或已断开")
            
        result = controller.swipe(data.x1, data.y1, data.x2, data.y2, data.duration)
        return {"status": "success", "message": result}
    except (DeviceConnectionError, CommandExecutionError) as e:
        raise HTTPException(status_code=400, detail=str(e))





@app.get("/devices")
def list_connected_devices():
    """列出所有已连接的设备"""
    try:
        stdout, stderr = subprocess.run(
            [ADB_PATH, 'devices'],
            capture_output=True,
            text=True,
            check=True
        ).stdout, subprocess.run(
            [ADB_PATH, 'devices'],
            capture_output=True,
            text=True,
            check=True
        ).stderr
        
        if stderr:
            raise CommandExecutionError(f"获取设备列表失败: {stderr}")
            
        lines = stdout.strip().split('\n')[1:]  # 跳过标题行
        connected_devices = []
        
        for line in lines:
            if line.strip():
                device_id, status = line.strip().split('\t', 1)
                if status.lower() == 'device':
                    connected_devices.append({
                        'id': device_id,
                        'status': status,
                        'managed_by_service': device_id in devices
                    })
                    
        return {"devices": connected_devices}
    except CommandExecutionError as e:
        raise HTTPException(status_code=400, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        # 改为导入字符串："模块名:应用实例名"
        "adb_http_server:app",  # 假设文件名为 match_http_server.py
        host="0.0.0.0",
        port=8001,
        reload=False,  # 开发模式启用热重载
        # workers=4,  # 生产环境按需开启
    )